from pocketflow import AsyncFlow
from typing import Optional, Dict

from .microflow import RetrieveAction, GenerateAction, Completed
from scorpio.core.common import get_logger

logger = get_logger(__name__)

class RAGAgent:
    """RAG Agent，使用pocketflow AsyncFlow实现文档处理流程"""
    
    shared: Dict
    _async_flow: Optional[AsyncFlow] = None

    def __init__(self, shared: Dict):
        self.name="rag"
        self.shared = shared
        """Get the search query from the shared store.
        shared ={
            "context": {
                "db_manager": db_manager,
                "processor": query_processor,
                "embedding_client": embedding_client,
                "generation": generation_service
            },
            "query": {
                "stream": True,
                "question": question
            }
        }
        """
        # 初始化pocketflow节点
        self.retrieve = RetrieveAction()
        self.generate = GenerateAction()
        self.completed = Completed() 
    def create_flow(self) -> AsyncFlow:
        """创建RAG流程"""
         # 连接流程节点
        self.retrieve - "failed" >> self.completed
        self.retrieve - "retrieved" >> self.generate
        
        self.generate - "generated" >> self.completed
        self.generate - "failure" >> self.completed
        
        # 创建并返回流程
        self._async_flow = AsyncFlow(start=self.retrieve)
        return self._async_flow
        
    
    async def run(self):
        logger.info("start to run flow: RAGFlow")
        return await self._async_flow.run_async(self.shared)
